


package graph

import (
    "container/vector";
    "container/heap";
    "fmt";
    
)





type WeightedEdge interface {
    Dest() Vertex;
    
    Source() Vertex;
    
    Weight() int64;
}

type WeightedGraph interface {
    
    OutEdges(n Vertex)  EdgeIterator;
    
    Edges() EdgeIterator;
    
    SortedUndirectedEdges() EdgeIterator;
    
    Vertices()  VertexIterator;
    
    NumVertices() int;
    
    Directed() bool;
    
    HasVertex(v Vertex) bool;
    
}


type VertexIterator interface {
    End() bool;
    Next() Vertex;
    Current() Vertex;
}

type EdgeIterator interface {
    End() bool;
    Next() WeightedEdge;
    Current() WeightedEdge;
}








func newHeap() *container {
    var h container;
    h.v = vector.New(0);
    h.pos = make(map[Vertex] int);
    return &h;
}


// Structure to use in Prim's and Dijkstra algorithms.
type nodeDist struct {
    node  Vertex;
    dist  int64;
}


func (n* nodeDist) Less(o interface{}) bool {
    return n.dist < o.(*nodeDist).dist;
}

func (n* nodeDist) String() string {
    return fmt.Sprintf("%s,%d", n.node, n.dist);
}



// container to use for the priority queue.
type container struct {
     v *vector.Vector;
     
     // position map to store positions of elements in the vector.
     // The following invariant is mantained.  if v.At(i) == A then pos[A] == i
     pos map[Vertex] int;
}


func (c *container) Push(x interface{}) {
    c.v.Push(x);
    c.pos[x.(*nodeDist).node] = (c.v.Len() - 1);
}


func (c *container) Pop() interface{} {
    val := c.v.Pop().(*nodeDist);
    c.pos[val.node] = 0, false;
    return val;
}

func (c *container) Swap(i, j int) {
    c.v.Swap(i, j);
    c.pos[c.v.At(j).(*nodeDist).node] = j;
    c.pos[c.v.At(i).(*nodeDist).node] = i;
}

func (c *container) Less(i, j int) bool {
    return c.v.Less(i, j);
}

func (c *container) Len() int {
    return c.v.Len();
}
//////////////////////////





func ExtractPath(start Vertex, t Vertex, distMap map[Vertex] int64, predecessorMap map[Vertex] Vertex) 
                (dist int64, path *vector.Vector) {
    _, present := predecessorMap[t];
    if !present {
        return INF, nil;
    }
    path = vector.New(0);
    path.Push(t);
    var p Vertex = t;
    for {
        if p == start {
            break;
        }
        p = predecessorMap[p];
        if p == nil {
            return INF, nil;
        }
        path.Push(p)
    }
    return distMap[t], path;
}


func  ShortestPathDijkstra(g WeightedGraph, s Vertex) (distMap map[Vertex] int64, predecessorMap map[Vertex] Vertex) {
    distMap  = make(map[Vertex] int64);
    colorMap := make(map[Vertex] Color);
    predecessorMap = make(map[Vertex] Vertex);
    vIter := g.Vertices();
    for !vIter.End() {
        v := vIter.Next();
        distMap[v] = INF;
        colorMap[v] = WHITE;
        predecessorMap[v] = nil;
    }
    distMap[s] = 0;
    predecessorMap[s] = s;
    var queue = newHeap();
    heap.Push(queue, &nodeDist{s, 0});
    for {
        if queue.Len() == 0 {
            break;
        }
        v := heap.Pop(queue).(*nodeDist);
        edgeIter := g.OutEdges(v.node);
        for !edgeIter.End() {
            var edge WeightedEdge  = edgeIter.Next();
            if colorMap[edge.Dest()] != BLACK {
                distFromV := distMap[v.node] + edge.Weight();
                if distMap[edge.Dest()] > distFromV {
                    predecessorMap[edge.Dest()] = v.node;
                    distMap[edge.Dest()] = distFromV;
                    if colorMap[edge.Dest()] == WHITE {
                        heap.Push(queue, &nodeDist{edge.Dest(), distFromV});
                        colorMap[edge.Dest()] = GREY;
                    } else if colorMap[edge.Dest()] == GREY {
                        heap.Remove(queue, queue.pos[edge.Dest()]);
                        heap.Push(queue, &nodeDist{edge.Dest(), distFromV});
                    }
                }
            }
        }
        colorMap[v.node] = BLACK;
    }
    return distMap, predecessorMap;
}





func  MinTreePrim(g WeightedGraph, start Vertex) (parent map[Vertex] Vertex, totalCost int64) {
    distMap  := make(map[Vertex] int64);
    colorMap := make(map[Vertex] Color);
    parent = make(map[Vertex] Vertex);
    vIter := g.Vertices();
    for !vIter.End() {
        v := vIter.Next();
        distMap[v] = INF;
        colorMap[v] = WHITE;
        parent[v] = nil;
    }
    distMap[start] = 0;
    parent[start] = start;
    var queue = newHeap();
    heap.Push(queue, &nodeDist{start, 0});
    totalCost = 0;
    for {
        if queue.Len() == 0 {
            break;
        }
        v := heap.Pop(queue).(*nodeDist);
        edgeIter := g.OutEdges(v.node);
        for !edgeIter.End() {
            var edge WeightedEdge  = edgeIter.Next();
            var u Vertex = edge.Dest();
            if colorMap[u] != BLACK {
                
                // fmt.Printf("Exporing edge %d %s\n", v.node, edge);
                if  distMap[u] > edge.Weight() {
                    parent[u] = v.node;
                    distMap[u] = edge.Weight();
                
                    if colorMap[u] == WHITE {
                        heap.Push(queue, &nodeDist{u, edge.Weight()});
                        colorMap[u] = GREY;
                    } else if colorMap[u] == GREY {
                        heap.Remove(queue, queue.pos[u]);
                        heap.Push(queue, &nodeDist{u, edge.Weight()});
                    }
                }
            }
            
        }
        totalCost += v.dist;
        colorMap[v.node] = BLACK;
    }
    return parent, totalCost;    
}


func  MinTreeKruskal(g WeightedGraph) (edges *vector.Vector, totalCost int64) {
    totalCost = 0;
    uf := NewUnionFind(g.Vertices());
    edgeIter := g.SortedUndirectedEdges();
    minTree := vector.New(0);
    for !edgeIter.End() {
        e := edgeIter.Next();
        if !uf.SameComponents(e.Source(), e.Dest()) {
            fmt.Printf("%s\n", e);
            uf.FindUnion(e.Source(), e.Dest());
            minTree.Push(e);
            totalCost += e.Weight();
        }
    }
    return minTree, totalCost;
}


    
    

func ShortestPathBellmanFord(g WeightedGraph, start Vertex) (distMap map[Vertex] int64, predecessorMap map[Vertex] Vertex) {
    if !g.HasVertex(start) {
        fmt.Printf("Graph does not have vertex %s\n", start);
        return nil, nil;
    }
    n := g.NumVertices();
    distMap = make(map[Vertex] int64);
    predecessorMap = make(map[Vertex] Vertex);
    vIter := g.Vertices();
    for !vIter.End() {
        v := vIter.Next();
        distMap[v] = INF;
        predecessorMap[v] = nil;
    }
    distMap[start] = 0;
    predecessorMap[start] = start;
    for i := 0; i < n; i++ {
        edgeIter := g.Edges();
        for !edgeIter.End() {
            e := edgeIter.Next();
            v := e.Source();
            u := e.Dest();
            if distMap[v] == INF {
                continue;
            }
            if distMap[u] > distMap[v] + e.Weight() {
                distMap[u] = distMap[v] + e.Weight();
                predecessorMap[u] = v;
            }
        }
    }
    
    edgeIter := g.Edges();
    for !edgeIter.End() {
        e := edgeIter.Next();
        v := e.Source();
        u := e.Dest();
        if distMap[u] > distMap[v] + e.Weight() {
            return nil, nil;
        }
    }
    
    return distMap, predecessorMap;
}





